import tushare as ts
import pandas as pd
import talib
from sklearn.ensemble import RandomForestClassifier
from sklearn.model_selection import train_test_split
from sklearn.metrics import accuracy_score, classification_report
import numpy as np
# 初始化pro接口
pro = ts.pro_api('1c7f85b9026518588c0d0cdac712c2d17344332c9c8cfe6bc83ee75c')
# 选择一支股票
ts_code = '000001.SZ'
# 获取3 - 5年行情数据，这里以3年为例
start_date = '20220415'
end_date = '20250415'
# 拉取数据
df = pro.daily(ts_code=ts_code, start_date=start_date, end_date=end_date, fields=[
    "ts_code",
    "trade_date",
    "open",
    "high",
    "low",
    "close",
    "pre_close",
    "change",
    "pct_chg",
    "vol",
    "amount"
])
# 对数据进行排序
df = df.sort_values(by='trade_date')
# 计算技术指标与特征值
df['ma5'] = talib.SMA(df['close'], timeperiod=5)
df['ma10'] = talib.SMA(df['close'], timeperiod=10)
df['rsi'] = talib.RSI(df['close'], timeperiod=14)
df['macd'], df['macdsignal'], df['macdhist'] = talib.MACD(df['close'], fastperiod=12, slowperiod=26, signalperiod=9)
# 计算收益率
df['return'] = df['close'].pct_change()
df = df.dropna()
# 创建目标变量，1表示上涨，0表示下跌
df['target'] = np.where(df['return'].shift(-1) > 0, 1, 0)
df = df.dropna()
# 准备特征和目标变量
features = ['open', 'high', 'low', 'close', 'vol', 'amount', 'ma5', 'ma10', 'rsi', 'macd', 'macdsignal', 'macdhist']
X = df[features]
y = df['target']
# 划分训练集和测试集
X_train, X_test, y_train, y_test = train_test_split(X, y, test_size=0.2, random_state=42)
# 建立随机森林模型
model = RandomForestClassifier(n_estimators=100, random_state=42)
model.fit(X_train, y_train)
# 预测
y_pred = model.predict(X_test)
# 模型评价
accuracy = accuracy_score(y_test, y_pred)
print(df[['ma5','ma10','rsi','macd','return']])
print(f"模型准确率: {accuracy}")
print("分类报告:")
print(classification_report(y_test, y_pred))
